from PyQt5.QtCore import QThread, pyqtSignal
from ultralytics import YOLO
import yaml
import os
import sys

class StreamToSignal:
    """把 print 输出直接发给 PyQt 信号"""
    def __init__(self, signal):
        self.signal = signal

    def write(self, msg):
        msg = msg.strip()
        if msg:
            self.signal.emit(msg)

    def flush(self):
        pass  # 必须有 flush 方法

class TrainThread(QThread):
    log_signal = pyqtSignal(str)  # 用于发送日志到界面

    def __init__(self, yaml_path=None):
        super().__init__()
        # 默认 train.yaml 在当前目录
        if yaml_path is None:
            self.yaml_path = os.path.join(os.path.dirname(__file__), "train.yaml")
        else:
            self.yaml_path = yaml_path
        self.is_running = True

    def run(self):
        # ---------------- 读取训练参数 ----------------
        try:
            with open(self.yaml_path, "r", encoding="utf-8") as f:
                params = yaml.safe_load(f) or {}
        except Exception as e:
            self.log_signal.emit(f"[错误] 读取 {self.yaml_path} 失败: {e}")
            return

        weights = params.get("weights", "yolo11n.pt")

        # ---------------- 处理 data.yaml 路径 ----------------
        default_data_yaml = os.path.join(os.path.dirname(__file__), "dataset/dog/dataset/data.yaml")
        data = params.get("data", default_data_yaml)
        data = os.path.abspath(data)

        # 确认文件存在
        if not os.path.isfile(data):
            self.log_signal.emit(f"[错误] 指定的 data.yaml 不存在 → {data}")
            return

        epochs = int(params.get("epochs", 100))
        batch = int(params.get("batch", 16))
        imgsz = int(params.get("imgsz", 640))
        workers = int(params.get("workers", 0))
        device = params.get("device", "cuda")

        self.log_signal.emit(f"[开始] 训练: weights={weights}, data={data}, epochs={epochs}")

        # ---------------- 启动训练 ----------------
        try:
            model = YOLO(weights)

            # 重定向 stdout，每行训练日志实时发送到界面
            sys.stdout = StreamToSignal(self.log_signal)

            model.train(
                data=data,
                epochs=epochs,
                imgsz=imgsz,
                batch=batch,
                workers=workers,
                device=device
            )

            # 恢复 stdout
            sys.stdout = sys.__stdout__
            self.log_signal.emit("[完成] 训练完成！")

        except Exception as e:
            sys.stdout = sys.__stdout__
            self.log_signal.emit(f"[错误] 训练出错: {repr(e)}")